package main

import (
	"encoding/json"
	"fmt"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
)

// ASTNode interface for serializing AST nodes
type ASTNode interface{}

// Custom type for serializing AST nodes
type Node struct {
	Type    string    `json:"type"`
	Pos     string    `json:"pos,omitempty"`
	Decls   []ASTNode `json:"decls,omitempty"`
	Imports []string  `json:"imports,omitempty"`
	Specs   []ASTNode `json:"specs,omitempty"`
	Stmts   []ASTNode `json:"stmts,omitempty"`
	Names   []string  `json:"names,omitempty"`
	Value   string    `json:"value,omitempty"` // Added to show values
	Func    ASTNode   `json:"func,omitempty"` // Added for function expressions
	Args    []ASTNode `json:"args,omitempty"` // Added for function arguments
}

// Convert Go AST node to a custom JSON-friendly structure
func nodeToJSON(node interface{}) ASTNode {
	switch n := node.(type) {
	case *ast.File:
		return Node{
			Type:    "File",
			Decls:   convertDecls(n.Decls),
			Imports: convertImports(n.Imports),
		}
	case *ast.GenDecl:
		return Node{
			Type:  "GenDecl",
			Specs: convertSpecs(n.Specs),
		}
	case *ast.FuncDecl:
		return Node{
			Type:  "FuncDecl",
			Stmts: convertStmts([]ast.Stmt{n.Body}),
		}
	case *ast.BlockStmt:
		return Node{
			Type:  "BlockStmt",
			Stmts: convertStmts(n.List),
		}
	case *ast.ExprStmt:
		return Node{
			Type:  "ExprStmt",
			Value: formatExpr(n.X),
		}
	case *ast.CallExpr:
		return Node{
			Type:  "CallExpr",
			Func:  formatExpr(n.Fun),
			Args:  convertExprs(n.Args),
		}
	case *ast.BasicLit:
		return Node{
			Type:  "BasicLit",
			Value: n.Value,
		}
	case *ast.Ident:
		return Node{
			Type:  "Ident",
			Value: n.Name,
		}
	case *ast.SelectorExpr:
		return Node{
			Type:  "SelectorExpr",
			Func:  formatExpr(n.X),
			Value: n.Sel.Name,
		}
	// Add cases for other AST types as needed
	default:
		return Node{Type: "Unknown"}
	}
}

func convertDecls(decls []ast.Decl) []ASTNode {
	nodes := make([]ASTNode, len(decls))
	for i, d := range decls {
		nodes[i] = nodeToJSON(d)
	}
	return nodes
}

func convertImports(imports []*ast.ImportSpec) []string {
	result := make([]string, len(imports))
	for i, imp := range imports {
		result[i] = imp.Path.Value
	}
	return result
}

func convertSpecs(specs []ast.Spec) []ASTNode {
	nodes := make([]ASTNode, len(specs))
	for i, s := range specs {
		nodes[i] = nodeToJSON(s)
	}
	return nodes
}

func convertStmts(stmts []ast.Stmt) []ASTNode {
	nodes := make([]ASTNode, len(stmts))
	for i, s := range stmts {
		nodes[i] = nodeToJSON(s)
	}
	return nodes
}

func convertExprs(exprs []ast.Expr) []ASTNode {
	nodes := make([]ASTNode, len(exprs))
	for i, e := range exprs {
		nodes[i] = nodeToJSON(e)
	}
	return nodes
}

// Format expressions for JSON output
func formatExpr(expr ast.Expr) string {
	switch e := expr.(type) {
	case *ast.Ident:
		return e.Name
	case *ast.BasicLit:
		return e.Value
	case *ast.SelectorExpr:
		return fmt.Sprintf("%s.%s", formatExpr(e.X), e.Sel.Name)
	case *ast.CallExpr:
		funcExpr := formatExpr(e.Fun)
		args := ""
		for i, arg := range e.Args {
			if i > 0 {
				args += ", "
			}
			args += formatExpr(arg)
		}
		return fmt.Sprintf("%s(%s)", funcExpr, args)
	default:
		return "Unknown"
	}
}

func main() {
	if len(os.Args) < 2 {
		fmt.Println("Usage: go run parser.go <file.go>")
		return
	}

	file := os.Args[1]
	fset := token.NewFileSet()
	node, err := parser.ParseFile(fset, file, nil, parser.AllErrors)
	if err != nil {
		fmt.Println("Error parsing file:", err)
		return
	}

	// Convert AST to JSON-friendly format
	jsonData, err := json.MarshalIndent(nodeToJSON(node), "", "  ")
	if err != nil {
		fmt.Println("Error converting AST to JSON:", err)
		return
	}

	fmt.Println(string(jsonData))
}
